package exploit

import (
	"github.com/cdk-team/CDK/pkg/plugin"
	"github.com/cdk-team/CDK/pkg/tool/ps"
	"log"
	"os"
	"runtime"
	"syscall"
)

type Event interface{}

type Tracee struct {
	proc   *os.Process
	events chan Event
	err    chan error

	cmds chan func()
}

func PtraceExec(name string, argv []string) (*Tracee, error) {
	t := &Tracee{
		events: make(chan Event, 1),
		err:    make(chan error, 1),
		cmds:   make(chan func()),
	}

	err := make(chan error)
	proc := make(chan *os.Process)
	go func() {
		runtime.LockOSThread()
		p, e := os.StartProcess(name, argv, &os.ProcAttr{
			Files: []*os.File{os.Stdin, os.Stdout, os.Stderr},
			Sys: &syscall.SysProcAttr{
				Ptrace: true,
			},
		})
		proc <- p
		err <- e
		if e != nil {
			return
		}
		go t.wait()
		t.trace()
	}()
	t.proc = <-proc
	return t, <-err
}

func (t *Tracee) wait() {
	defer close(t.events)
	for {
		state, err := t.proc.Wait()
		if err != nil {
			t.err <- err
			return
		}
		if state.Exited() {
			t.events <- Event(state.Sys().(syscall.WaitStatus))
			return
		}
		t.events <- Event(state.Sys().(syscall.WaitStatus))
	}
}

func (t *Tracee) trace() {
	for cmd := range t.cmds {
		cmd()
	}
}

func CheckPidInject() bool {
	cmd := "/bin/echo"
	arg := []string{
		"1",
	}
	t, err := PtraceExec(cmd, arg)
	if err != nil {
		log.Println("exec ptrace failed. err:", err)
		return false
	} else {
		log.Println("exec ptrace success.", t.events)
		ps.RunPs()
		return true
	}
}

// plugin interface
type PidInject struct{}

func (p PidInject) Desc() string {
	return "check if pid injection works with cap=SYS_PTRACE. usage: ./cdk run check-ptrace"
}
func (p PidInject) Run() bool {
	CheckPidInject()
	return true
}

func init() {
	exploit := PidInject{}
	plugin.RegisterExploit("check-ptrace", exploit)
}
